package middleware

import (
	"bytes"
	"io"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

// LoggerConfig 日志配置
type LoggerConfig struct {
	// 跳过日志记录的路径
	SkipPaths []string
	// 是否记录请求体
	LogRequestBody bool
	// 是否记录响应体
	LogResponseBody bool
	// 请求体最大记录长度
	MaxRequestBodySize int
	// 响应体最大记录长度
	MaxResponseBodySize int
}

// DefaultLoggerConfig 默认日志配置
var DefaultLoggerConfig = LoggerConfig{
	SkipPaths:           []string{"/health", "/metrics", "/favicon.ico"},
	LogRequestBody:      false,
	LogResponseBody:     false,
	MaxRequestBodySize:  1024, // 1KB
	MaxResponseBodySize: 1024, // 1KB
}

// Logger 日志中间件
func Logger(config ...LoggerConfig) gin.HandlerFunc {
	// 使用默认配置或传入的配置
	loggerConfig := DefaultLoggerConfig
	if len(config) > 0 {
		loggerConfig = config[0]
	}

	return func(c *gin.Context) {
		// 如果路径在跳过列表中，则直接跳过
		path := c.Request.URL.Path
		for _, skipPath := range loggerConfig.SkipPaths {
			if path == skipPath {
				c.Next()
				return
			}
		}

		// 开始时间
		start := time.Now()

		// 请求方法
		method := c.Request.Method

		// 请求路径
		reqPath := c.Request.URL.Path

		// 请求IP
		clientIP := c.ClientIP()

		// 请求头
		headers := c.Request.Header

		// 请求参数
		query := c.Request.URL.RawQuery

		// 记录请求体
		var requestBody []byte
		if loggerConfig.LogRequestBody && c.Request.Body != nil {
			// 读取请求体
			body, err := io.ReadAll(c.Request.Body)
			if err == nil {
				// 限制请求体大小
				if len(body) > loggerConfig.MaxRequestBodySize {
					requestBody = body[:loggerConfig.MaxRequestBodySize]
				} else {
					requestBody = body
				}
				// 重新设置请求体，以便后续处理
				c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
			}
		}

		// 捕获响应体
		if loggerConfig.LogResponseBody {
			// 创建自定义ResponseWriter
			blw := &bodyLogWriter{
				ResponseWriter: c.Writer,
				body:           bytes.NewBufferString(""),
			}
			c.Writer = blw

			// 处理请求
			c.Next()

			// 获取响应体
			responseBody := blw.body.Bytes()

			// 限制响应体大小
			if len(responseBody) > loggerConfig.MaxResponseBodySize {
				responseBody = responseBody[:loggerConfig.MaxResponseBodySize]
			}

			// 记录请求和响应信息
			zap.L().Info("HTTP请求",
				zap.String("method", method),
				zap.String("path", reqPath),
				zap.String("query", query),
				zap.String("ip", clientIP),
				zap.Any("headers", headers),
				zap.Int("status", c.Writer.Status()),
				zap.Int("size", c.Writer.Size()),
				zap.Duration("latency", time.Since(start)),
				zap.ByteString("request", requestBody),
				zap.ByteString("response", responseBody),
			)
		} else {
			// 处理请求
			c.Next()

			// 记录请求信息
			zap.L().Info("HTTP请求",
				zap.String("method", method),
				zap.String("path", reqPath),
				zap.String("query", query),
				zap.String("ip", clientIP),
				zap.Any("headers", headers),
				zap.Int("status", c.Writer.Status()),
				zap.Int("size", c.Writer.Size()),
				zap.Duration("latency", time.Since(start)),
				zap.ByteString("request", requestBody),
			)
		}
	}
}

// bodyLogWriter 用于捕获响应体的自定义ResponseWriter
type bodyLogWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

// Write 重写Write方法，同时写入响应体和缓冲区
func (w *bodyLogWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

// WriteString 重写WriteString方法，同时写入响应体和缓冲区
func (w *bodyLogWriter) WriteString(s string) (int, error) {
	w.body.WriteString(s)
	return w.ResponseWriter.WriteString(s)
}
